import torch
import torchvision.datasets
from torch import nn
from torch.utils.data import DataLoader

dateset = torchvision.datasets.CIFAR10("../../data", train=False,
                                       transform=torchvision.transforms.ToTensor(), download=True)

dataLoader = DataLoader(dateset,batch_size=64)

class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.linear1 = nn.Linear(196608,10)

    def forward(self,x):
        x = self.linear1(x)
        return x


module = Model()

for data in dataLoader:
    imgs,targets = data
    print(imgs.shape)
    output = torch.flatten(imgs)
    print(output.shape)
    output = module(output)
    print(output.shape)
    break